# -*- coding: utf-8 -*-
"""
Created on Sat Jun 23 12:14:27 2018

@author: 301
"""

import sim
import os
import time
import jieba
import pandas as pd


def text2set(text, cut_all):
    l = jieba.cut(text, cut_all=cut_all)
    s = set([x for x in l if sim.chn.match(x) and x not in sim.stop_list])
    return s


if __name__ == "__main__":

    path = input("请输入待检索文件目录：")
    file_list = [
        x for x in os.listdir(path) if os.path.splitext(x)[-1] == '.txt'
    ]
    print("{:}个待检索文件读取成功!".format(len(file_list)))

    #csvpath = input("input csv file path:")
    csvpath = r'C:\Users\luther\Desktop\\file_list.csv'
    df = pd.read_csv(csvpath)
    s = []
    s_full = []

    ret_path = input("input output file path:")

    start_time = time.time()
    for i in range(len(df)):
        code = (6 - len(str(int(df.stkcd[i])))) * '0' + str(int(df.stkcd[i]))
        year = df.year[i]
        file = '{}-{}.txt'.format(year, code)
        text = sim.readtext(path, file)
        ws = text2set(text, False)
        ws_full = text2set(text, True)
        s.append(ws)
        s_full.append(ws_full)
        end_time = time.time()
        remain_time = (end_time - start_time) / (i + 1) * (len(df) - (i + 1))
        print(
            "\r{:>6.2f}% Done, Time Remained: {:04}:{:02}:{:02}".format(
                (i + 1) * 100 / len(df), int(remain_time // 3600),
                int(remain_time // 60 % 60), int(remain_time % 60)),
            end="")

    df['WordSet'] = s
    df['WordSet_Full'] = s_full
    df.to_csv(ret_path, index=False, encoding='utf-8')
